import os


def doit(dir):   # 得到种类类别数量
    alls = {}
    labes = 0
    for filename in os.listdir(dir):  # 遍历目录下所有文件
        lens = len(filename)
        if filename[-4:] == ".txt":
            continue
        for i in range(1, lens + 1):
            if filename[-i] == '_':
                name = filename[:lens - i]
                if name not in alls.keys():
                     alls[name] = labes
                     labes += 1

                break
    return alls


def saveClasstoTxt(dir):
    file_path = "D:/work/test/classes.txt"
    LABELS = doit(dir)
    with open(file_path, "w") as f:
        for item in LABELS.keys():
            if not item.islower():
                f.write("%s\n" % item)

    print("File saved to:", os.path.abspath(file_path))